import socket
import time
import struct
from PIL import Image
from io import BytesIO
import random

NIL = 0
ACK = 1
NAK = 2
ACK0 = 3
ACK1 = 4

VER = '3.0'

class RDT30Server:
    def __init__(self, addr : tuple[str, int], errorsize: int) -> None:
        self.serverSocket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.serverSocket.bind(addr)
        self.data = bytearray()
        self.tmpdata = None
        self.target = None

        # 0 -> waiting for 0 from below
        # 1 -> waiting for 1 from below
        self.state = 0
        self.errorsize = errorsize
        pass


    def recvpkt(self) -> bytes:
        data, target = self.serverSocket.recvfrom(1024)
        self.target = target
        return data
    
    def makepkt(self, rawdata: bytes, confirmation: int, sequence: int):
        # rdt 2.1 报文格式:
        # ack/nak 1byte
        # sequence 1byte
        # length 2byte
        # checksum 2byte
        pkt = struct.pack(
            "!BBHH%ds" % len(rawdata), confirmation, sequence, len(rawdata), 0, rawdata
        )
        csum = self.checksum(pkt)
        return struct.pack(
            "!BBHH%ds" % len(rawdata), confirmation, sequence, len(rawdata), csum, rawdata
        )

    def sendpkt(self, pkt: bytes):
        self.serverSocket.sendto(pkt, self.target)

    def checksum(self, data: bytes) -> int:
        csum = 0
        for i in range(1, len(data), 2):
            csum += int.from_bytes(data[i - 1 : i + 1], "big")
        if len(data) % 2 == 1:
            csum += data[-1] << 8 & 0xFF00
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = ~csum & 0xFFFF
        return csum
    
    def valid(self, pkt: bytes) -> bool:
        # print('[dbg] ', self.checksum(pkt))
        return self.checksum(pkt) == 0

    def extract(self, pkt: bytes) -> tuple[int, int, int, int, bytes]:
        confirmation, sequence, length, csum = struct.unpack("!BBHH", pkt[:6])
        data = struct.unpack_from("!%ds" % length, pkt, 6)[0]
        return confirmation, sequence, length, csum, data
    
    def biterror(self, pkt: bytes):
        if random.randint(1, self.errorsize) == 1:
            pos = random.randint(0, len(pkt) - 1)
            off = random.randint(0, 7)
            npkt = bytearray(pkt)
            npkt[pos] ^= 1 << off
            return bytes(npkt)
        return pkt

    def run(self):
        while True:

            match self.state:
                case 0:
                    rcvdpkt = self.recvpkt()
                    if self.valid(rcvdpkt) and self.extract(rcvdpkt)[1] == 0:
                        confirmation, sequence, length, csum, self.tmpdata = self.extract(rcvdpkt)
                        if length == 0:
                            print('data fully received.')
                            self.sendpkt(self.makepkt(b'', ACK0, 0))
                            break
                        print('sq 0 data received at ', len(self.data))
                        self.data += self.tmpdata
                        self.sendpkt(self.makepkt(b'', ACK0, 0))
                        self.state = 1
                    else :
                            self.sendpkt(self.biterror(self.makepkt(b'', ACK1, 0)))
                            self.state = 0
                case 1:
                    rcvdpkt = self.recvpkt()
                    if self.valid(rcvdpkt) and self.extract(rcvdpkt)[1] == 1:
                        confirmation, sequence, length, csum, self.tmpdata = self.extract(rcvdpkt)
                        if length == 0:
                            print('data fully received.')
                            self.sendpkt(self.makepkt(b'', ACK1, 0))
                            break
                        print('sq 1 data received at ', len(self.data))
                        self.data += self.tmpdata
                        self.sendpkt(self.makepkt(b'', ACK1, 0))
                        self.state = 0
                    else :
                            self.sendpkt(self.biterror(self.makepkt(b'', ACK0, 0)))
                            self.state = 1

                    # if self.valid(rcvdpkt):
                    #     confirmation, sequence, length, csum, self.tmpdata = self.extract(rcvdpkt)
                    #     if sequence == 1:
                    #         self.sendpkt(self.biterror(self.makepkt(b'', ACK, 0)))
                    #         self.state = 0
                    #     else:
                    #         if length == 0:
                    #             print('data fully received.')
                    #             self.sendpkt(self.makepkt(b'', ACK, 0))
                    #             break
                    #         print('data received at ', len(self.data))
                    #         self.data += self.tmpdata
                    #         self.sendpkt(self.makepkt(b'', ACK, 0))
                    #         self.state = 1

                    # else:
                    #     self.sendpkt(self.biterror(self.makepkt(b'', NAK, 0)))
                    #     self.state = 0
                # case 1:
                #     rcvdpkt = self.recvpkt()
                #     if self.valid(rcvdpkt):
                #         confirmation, sequence, length, csum, self.tmpdata = self.extract(rcvdpkt)
                #         if sequence == 0:
                #             self.sendpkt(self.biterror(self.makepkt(b'', ACK, 1)))
                #             self.state = 1
                #         else:
                #             if length == 0:
                #                 print('data fully received.')
                #                 self.sendpkt(self.makepkt(b'', ACK, 0))
                #                 break
                #             print('data received at ', len(self.data))
                #             self.data += self.tmpdata
                #             self.sendpkt(self.makepkt(b'', ACK, 0))
                #             self.state = 0

                #     else:
                #         self.sendpkt(self.biterror(self.makepkt(b'', NAK, 1)))
                #         self.state = 1

            pass
        pass

s = RDT30Server(('127.0.0.1', 65432), 100)
s.run()

Image.open(BytesIO(s.data)).show()

open(VER + '-server.png', 'wb').write(s.data)

input("Press any key to exit...")
